#include <stdlib.h>
#include <unistd.h>
#include <signal.h>
#include <sys/time.h>

// 定义令牌桶结构
struct tbf_st {
    int token;    // 当前令牌数
    int cps;      // 每秒产生的令牌数
    int burst;    // 令牌桶的容量
};

// 定义全局变量
static struct tbf_st *jobs[1024] = {NULL};

// 初始化令牌桶
struct tbf_st *tbf_init(int cps, int burst) {
    struct tbf_st *tb = malloc(sizeof(struct tbf_st));
    if (tb == NULL) {
        return NULL;
    }
    tb->token = 0;
    tb->cps = cps;
    tb->burst = burst;
    return tb;
}

// 释放销毁令牌桶
void tbf_destroy(struct tbf_st *tb) {
    if (tb != NULL) {
        free(tb);
    }
}

// 取令牌
int tbf_fetch_token(struct tbf_st *tb) {
    if (tb == NULL) {
        return -1; // 无效的令牌桶
    }
    if (tb->token < tb->burst) {
        tb->token++;
        return 0; // 成功取得令牌
    } else {
        return -1; // 令牌桶已满
    }
}

// 信号处理函数，用于积攒令牌
static void alarm_handler(int s) {
    for (int i = 0; i < 1024; i++) {
        if (jobs[i] != NULL) {
            jobs[i]->token += jobs[i]->cps;
            if (jobs[i]->token > jobs[i]->burst) {
                jobs[i]->token = jobs[i]->burst;
            }
        }
    }
    alarm(1); // 每秒触发一次
}

// 初始化信号处理函数
void tbf_init_signal(void) {
    struct sigaction act;
    act.sa_handler = alarm_handler;
    sigemptyset(&act.sa_mask);
    act.sa_flags = 0;
    sigaction(SIGALRM, &act, NULL);
    alarm(1); // 第一次触发
}
